{-# OPTIONS -fglasgow-exts #-}

-----------------------------------------------------------------------------
-- |
-- Maintainer  :  joost.visser@di.uminho.pt, alcino@di.uminho.pt
-- Stability   :  experimental
-- Portability :  portable
--
-- Point-free Type-preserving Rewriting
--
-----------------------------------------------------------------------------

module Data.Transform.Rewriting where

import Data.Transform.Type

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.RWS
import Data.List

rmTags :: Type a -> Type a
rmTags (List a) = List (rmTags a)
rmTags (Set a) = Set (rmTags a)
rmTags (Map a b) = Map (rmTags a) (rmTags b)
rmTags (Either a b) = Either (rmTags a) (rmTags b)
rmTags (Prod a b) = Prod (rmTags a) (rmTags b)
rmTags (Func a b) = Func (rmTags a) (rmTags b)
rmTags (Tag _ a) = rmTags a
rmTags (Mu f) = Mu f      -- should descend into f as well...
rmTags (PF a) = PF (rmTags a)
rmTags x = x    



-- Generic queries

type GenericQ r = forall a. Type a -> a -> r

gmapQ :: GenericQ r -> GenericQ [r]
gmapQ q t x = aux q (toSpine t x)
    where aux :: GenericQ r -> (forall a. Spine a -> [r])
          aux q (_ `As` _)      = []
          aux q (Ap f (t :| x)) = aux q f ++ [q t x]

type Query r = forall a . Typed a -> r

everything :: (r -> r -> r) -> GenericQ r -> Query r
everything op q (t :| x) = aux op q t x
    where aux :: (r -> r -> r) -> GenericQ r -> GenericQ r
	  aux op q t x = foldl1 op ([q t x] ++ gmapQ (aux op q) t x)
	       
-- Locations and contexts

type Location = [Int]

down :: Location -> Location
down = (0:)

next :: Location -> Location
next (h:t) = (h+1):t

replace :: Location -> Typed b -> Typed a -> Maybe a
replace [0] (b :| x) (a :| y) = coerce b a x
replace l   (b :| x) (a :| y) = do s <- aux (last l) (init l) (toSpine a y)
                                   return $ fromSpine s
    where aux :: Int -> Location -> Spine a -> Maybe (Spine a)
          aux 0     l (Ap f (a :| y)) = do z <- replace l (b :| x) (a :| y)
                                           return $ Ap f (a :| z)
          aux (n+1) l (Ap f (a :| y)) = do g <- aux n l f
                                           return $ Ap g (a :| y)

hole :: Type a -> a
hole (PF _) = HOLE

puthole :: Location -> Dynamic -> Dynamic
puthole l (Dyn t x) = Dyn t (xua l (t :| x))
    where xua :: Location -> Typed a -> a
          xua [0] (a :| y) = hole a
	  xua l   (a :| y) = fromSpine $ aux (last l) (init l) (toSpine a y)
	  aux :: Int -> Location -> Spine a -> Spine a
          aux 0     l (Ap f (a :| y)) = Ap f (a :| xua l (a :| y))
          aux (n+1) l (Ap f (a :| y)) = Ap (aux n l f) (a :| y)

-- The basic type of rules
type GenericM m = forall a . Type a -> PF a -> m (PF a)

-- Generic one-level map with location update

gmapMo :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
gmapMo h t y = do s <- aux h (toSpine (PF t) y)
		  return (fromSpine s)
    where aux :: (MonadReader Location m, MonadPlus m) => GenericM m -> (Spine a -> m (Spine a))
          aux h (c `As` _) = mzero
          aux h (Ap f (PF t :| x)) = (do g <- local next $ aux h f
				         return (Ap g (PF t :| x)))
				     `mplus`
				     (do y <- local down $ h t x
				         return (Ap f (PF t :| y)))

gmapMo' :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
gmapMo' h t y = do s <- aux h (toSpine (PF t) y)
                   return (fromSpine s)
    where aux :: (MonadReader Location m, MonadPlus m) => GenericM m -> (Spine a -> m (Spine a))
          aux h (c `As` _) = mzero
          aux h (Ap f (PF t :| x)) = (do y <- local down $ h t x
                                         return (Ap f (PF t :| y)))
                                    `mplus`
                                     (do g <- local next $ aux h f
                                         return (Ap g (PF t :| x)))



gmapM :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
gmapM h t y = do s <- aux h (toSpine (PF t) y)
		 return (fromSpine s)
    where aux :: (MonadReader Location m, MonadPlus m) => GenericM m -> (Spine a -> m (Spine a))
          aux h x@(c `As` _) = return x
          aux h (Ap f (PF t :| x)) = do 
            g <- local next $ aux h f
            y <- local down $ h t x
            return (Ap g (PF t :| y))

-- Promoting rules to traversals by updating context and propagating the type
top :: (MonadReader Location m, MonadState Dynamic m, MonadPlus m) => GenericM m -> GenericM m
top f t x = do y <- f t x
	       (Dyn u c) <- get
	       l <- ask
	       let Just z = replace l (PF t :| y) (u :| c) 
	       put (Dyn u z)
	       return y


-- Strategy combinators
once :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
once f = f ||| gmapMo (once f)

onceBU :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
onceBU f = (gmapMo' (once f)) ||| f

everywhere :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
everywhere f = f >>> gmapM (everywhere f)

(>>>) :: Monad m => GenericM m -> GenericM m -> GenericM m
(f >>> g) t x = f t x >>= g t

(|||) :: MonadPlus m => GenericM m -> GenericM m -> GenericM m
(f ||| g) t x = f t x `mplus` g t x

many :: MonadPlus m => GenericM m -> GenericM m
many r = (r >>> many r) ||| nop

many1 :: MonadPlus m => GenericM m -> GenericM m
many1 r = r >>> many r

nop :: Monad m => GenericM m
nop t = return

try :: MonadPlus m => GenericM m -> GenericM m
try x = x ||| nop

-- Rewriting monad

type Rewrite = RWST Location [(String,String)] Dynamic Maybe

type Rule = GenericM Rewrite

success :: String -> a -> Rewrite a
success n x = 
    do z <- get
       tell [(n, gshow Dynamic z)]
       return x

context :: Rewrite (Typed Dynamic)
context =
    do l <- ask
       y <- get
       return (Dynamic :| puthole l y)



-- Simplification wrapers

simplify :: Typeable a => Bool -> Rule -> PF a -> IO (PF a)
simplify b s e = do
  Just (x,l) <- return $ evalRWST (s typeof e) [0] (Dyn typeof e)
  when b $ sequence_ (map aux l)
  putStrLn ("  "++(gshow typeof x))
  return x
 where aux (n,y) =putStrLn ("  "++y) >> putStrLn ("=   { "++n++" }")

reduceMaybe :: Rule -> Type a -> PF a -> Maybe (PF a)
reduceMaybe s t x = fmap fst (evalRWST (s t x) [0] (Dyn (PF t) x))

reduceIO :: Rule -> Type a -> PF a -> IO (PF a)
reduceIO s t x = case (evalRWST (s t x) [0] (Dyn (PF t) x)) of
  Nothing      -> fail "Reduction failed"
  Just (y,_)   -> return y 

reduce :: Rule -> Type a -> PF a -> PF a
reduce s t x = maybe x fst (evalRWST (s t x) [0] (Dyn (PF t) x))

reduce2 :: Typeable a => Rule -> PF a -> PF a
reduce2 s x = maybe x fst (evalRWST (s typeof x) [0] (Dyn typeof x))

ruleWithRmConstraints :: Rule -> Rule
ruleWithRmConstraints r t x = r (rmConstraints t) x


